import asyncio
from typing import List, Union

from lagent.llms.base_llm import AsyncBaseLLM, BaseLLM
from lagent.utils.util import filter_suffix


def asdict_completion(output):
    """
    将输出结果转换为字典格式
    
    参数:
        output: 模型输出对象
        
    返回:
        dict: 包含输出文本和相关信息的字典
    """
    return {
        key: getattr(output, key)
        for key in [
            'text',              # 生成的文本
            'token_ids',         # token ID序列
            'cumulative_logprob', # 累积对数概率
            'logprobs',          # 每个token的对数概率
            'finish_reason',     # 生成完成的原因
            'stop_reason'        # 停止生成的原因
        ]
    }


class VllmModel(BaseLLM):
    """
    vLLM模型的包装类
    
    提供了对vLLM模型的高级封装，支持模型加载和文本生成功能
    
    参数:
        path (str): 模型路径，可以是：
            - 本地huggingface模型目录路径
            - huggingface.co上的模型ID，如：
              "internlm/internlm-chat-7b"
              "Qwen/Qwen-7B-Chat"
              "baichuan-inc/Baichuan2-7B-Chat"
        tp (int): tensor并行度
        vllm_cfg (dict): vLLM模型初始化的其他参数
    """

    def __init__(self, path: str, tp: int = 1, vllm_cfg=dict(), **kwargs):
        super().__init__(path=path, **kwargs)
        from vllm import LLM
        # 初始化vLLM模型
        self.model = LLM(
            model=self.path,
            trust_remote_code=True,  # 信任远程代码
            tensor_parallel_size=tp,  # 设置tensor并行度
            **vllm_cfg
        )

    def generate(self,
                 inputs: Union[str, List[str]],
                 do_preprocess: bool = None,
                 skip_special_tokens: bool = False,
                 return_dict: bool = False,
                 **kwargs):
        """
        非流式模式下的文本生成
        
        参数:
            inputs (Union[str, List[str]]): 输入文本或文本列表
            do_preprocess (bool): 是否进行预处理，默认为True时会应用chat_template
            skip_special_tokens (bool): 是否移除特殊token，默认False
            return_dict (bool): 是否返回字典格式的结果
            
        返回:
            生成的文本或文本列表/字典格式的完整输出
        """
        from vllm import SamplingParams

        # 处理单条输入的情况
        batched = True
        if isinstance(inputs, str):
            inputs = [inputs]
            batched = False
        prompt = inputs
        
        # 更新生成参数
        gen_params = self.update_gen_params(**kwargs)
        max_new_tokens = gen_params.pop('max_new_tokens')
        stop_words = gen_params.pop('stop_words')

        # 配置采样参数
        sampling_config = SamplingParams(
            skip_special_tokens=skip_special_tokens,
            max_tokens=max_new_tokens,
            stop=stop_words,
            **gen_params
        )
        
        # 生成文本
        response = self.model.generate(prompt, sampling_params=sampling_config)
        texts = [resp.outputs[0].text for resp in response]
        # 移除停止词
        texts = filter_suffix(texts, self.gen_params.get('stop_words'))
        for resp, text in zip(response, texts):
            resp.outputs[0].text = text
            
        # 根据返回类型格式化输出
        if batched:
            return [asdict_completion(resp.outputs[0])
                    for resp in response] if return_dict else texts
        return asdict_completion(
            response[0].outputs[0]) if return_dict else texts[0]


class AsyncVllmModel(AsyncBaseLLM):
    """
    vLLM模型的异步包装类
    
    提供了对vLLM模型的异步操作支持，适用于高并发场景
    
    参数:
        path (str): 模型路径，可以是：
            - 本地huggingface模型目录路径
            - huggingface.co上的模型ID
        tp (int): tensor并行度
        vllm_cfg (dict): vLLM模型初始化的其他参数
    """

    def __init__(self, path: str, tp: int = 1, vllm_cfg=dict(), **kwargs):
        super().__init__(path=path, **kwargs)
        from vllm import AsyncEngineArgs, AsyncLLMEngine

        # 配置异步引擎参数
        engine_args = AsyncEngineArgs(
            model=self.path,
            trust_remote_code=True,
            tensor_parallel_size=tp,
            **vllm_cfg
        )
        # 初始化异步LLM引擎
        self.model = AsyncLLMEngine.from_engine_args(engine_args)

    async def generate(self,
                       inputs: Union[str, List[str]],
                       session_ids: Union[int, List[int]] = None,
                       do_preprocess: bool = None,
                       skip_special_tokens: bool = False,
                       return_dict: bool = False,
                       **kwargs):
        """
        异步模式下的文本生成
        
        参数:
            inputs: 输入文本或文本列表
            session_ids: 会话ID或ID列表，用于区分不同会话
            do_preprocess: 是否进行预处理
            skip_special_tokens: 是否移除特殊token
            return_dict: 是否返回字典格式的结果
            
        返回:
            生成的文本或文本列表/字典格式的完整输出
        """
        from vllm import SamplingParams

        # 处理输入格式
        batched = True
        if isinstance(inputs, str):
            inputs = [inputs]
            batched = False
        if session_ids is None:
            session_ids = list(range(len(inputs)))
        elif isinstance(session_ids, (int, str)):
            session_ids = [session_ids]
        assert len(inputs) == len(session_ids)

        prompt = inputs
        # 更新生成参数
        gen_params = self.update_gen_params(**kwargs)
        max_new_tokens = gen_params.pop('max_new_tokens')
        stop_words = gen_params.pop('stop_words')

        # 配置采样参数
        sampling_config = SamplingParams(
            skip_special_tokens=skip_special_tokens,
            max_tokens=max_new_tokens,
            stop=stop_words,
            **gen_params
        )

        # 定义异步生成函数
        async def _inner_generate(uid, text):
            resp, generator = '', self.model.generate(
                text, sampling_params=sampling_config, request_id=uid)
            async for out in generator:
                resp = out.outputs[0]
            return resp

        # 并行执行生成任务
        response = await asyncio.gather(*[
            _inner_generate(sid, inp) for sid, inp in zip(session_ids, prompt)
        ])
        texts = [resp.text for resp in response]
        # 移除停止词
        texts = filter_suffix(texts, self.gen_params.get('stop_words'))
        for resp, text in zip(response, texts):
            resp.text = text
            
        # 根据返回类型格式化输出
        if batched:
            return [asdict_completion(resp)
                    for resp in response] if return_dict else texts
        return asdict_completion(response[0]) if return_dict else texts[0]
